#include "cem/assembler.h"
#include "cem/boundary_condition_object.h"
#include "cem/electromagnetics_system.h"
#include "cem/material_database.h"
#include "cem/material_utils.h"

#include <libmesh/boundary_info.h>
#include <libmesh/dense_matrix.h>
#include <libmesh/dense_vector.h>
#include <libmesh/dof_map.h>
#include <libmesh/elem.h>
#include <libmesh/enum_order.h>
#include <libmesh/equation_systems.h>
#include <libmesh/fe.h>
#include <libmesh/linear_implicit_system.h>
#include <libmesh/mesh.h>
#include <libmesh/mesh_base.h>
#include <libmesh/numeric_vector.h>
#include <libmesh/parallel.h>
#include <libmesh/quadrature_gauss.h>
#include <libmesh/sparse_matrix.h>
#include <libmesh/vector_value.h>

#include <algorithm>
#include <cmath>
#include <complex>
#include <unordered_map>
#include <unordered_set>

namespace cem
{

void
Assembler::assemble(ElectromagneticsSystem & system)
{
    using namespace libMesh;

    auto & es = system.equation_systems();
    auto & sys = system.linear_system();

    // 先获取通信器以便打印rank
    const MeshBase & mesh = es.get_mesh();
    const unsigned int dim = mesh.mesh_dimension();

    // 准备全局对象并记录日志
    sys.matrix->zero();
    sys.rhs->zero();

    // 空网格短路保护：无活动单元或维度为0时不进行 FE 构建，直接清零并返回
    // 这样可避免在 NEDELEC 向量元上于 dim=0 的情况下构建 FE 导致的异常。
    if (mesh.n_active_elem() == 0 || dim == 0)
    {
        sys.matrix->close();
        sys.rhs->close();
        return;
    }

    const DofMap & dof_map = sys.get_dof_map();
    const unsigned int var = system.var_index();
    const FEType fe_type = sys.variable_type(var);

    // 有限元对象与积分规则
    std::unique_ptr<FEVectorBase> fe = FEVectorBase::build(dim, fe_type);
    QGauss qrule(dim, libMesh::SECOND);
    fe->attach_quadrature_rule(&qrule);
    // 面积分所需的 FE 与积分规则
    std::unique_ptr<FEVectorBase> fe_face = FEVectorBase::build(dim, fe_type);
    QGauss qface(dim > 0 ? dim - 1 : 0, SECOND);
    fe_face->attach_quadrature_rule(&qface);

    const std::vector<Real> & JxW = fe->get_JxW();
    // HCurl: 使用 curl 及向量形函数
    const std::vector<std::vector<RealVectorValue>> & curl = fe->get_curl_phi();
    const std::vector<std::vector<RealVectorValue>> & phi_v = fe->get_phi();

    // 局部矩阵/向量（复数 n×n 与 n）
    DenseMatrix<Number> ke;        // n × n（Number 可能为复数）
    DenseVector<Number> fe_local;  // n
    std::vector<dof_id_type> dofs; // 大小 n

    const double omega = system.omega();

    // 集中式材料管理器：已在 ElectromagneticsSystem::initialize 中完成初始化与缓存
    const MaterialDatabase & material_manager = system.material_manager();

    // 将边界规格整理为 id->spec 的快速查询表
    std::unordered_map<int, cem::BoundarySpec> bnd_spec_map;
    for (const auto & sp : system.boundary_specs())
        bnd_spec_map[sp.boundary_id] = sp;

    // 强施加Dirichlet所需：收集需钉扎的全局行索引
    std::unordered_set<libMesh::dof_id_type> dirichlet_rows_keep; // zero_rows(diag=1)
    // 非齐次切向 Dirichlet 的 L2 投影累计（按“行 id”累计复数 numer 与公共分母）
    std::unordered_map<libMesh::dof_id_type, std::complex<double>> dirichlet_value_num;
    std::unordered_map<libMesh::dof_id_type, double> dirichlet_value_den;

    // 遍历本地活动单元
    for (auto it = mesh.active_local_elements_begin(); it != mesh.active_local_elements_end(); ++it)
    {
        const Elem * elem = *it;

        // 该单元的自由度索引（复数变量）
        dof_map.dof_indices(elem, dofs, var);
        const unsigned int n = dofs.size();

        // 重置局部矩阵/向量尺寸
        ke.resize(n, n);
        fe_local.resize(n);
        ke.zero();
        fe_local.zero();
        // 绑定当前单元并初始化 FE（体积分）
        fe->reinit(elem);

        // 使用集中式材料管理器获取材料张量（σ, ε, μ^{-1}）
        const int sub_id = static_cast<int>(elem->subdomain_id());
        const auto & inv_mu = material_manager.get_inv_mu_tensor(sub_id);
        const auto & sigma_tensor = material_manager.get_sigma_tensor(sub_id);
        const auto & epsl_tensor = material_manager.get_epsl_tensor(sub_id);

        // 积分点循环（体项装配）
        const unsigned int n_qp = qrule.n_points();
        for (unsigned int qp = 0; qp < n_qp; ++qp)
        {
            for (unsigned int i = 0; i < n; ++i)
            {
                for (unsigned int j = 0; j < n; ++j)
                {
                    // 体项离散映射（复数形式）
                    // 方程: curl(μ^{-1} curl E) - i·ω·σ·E - ω²·ε·E = 0
                    // 矩阵条目: K - ω² M_ε - i·ω M_σ
                    const double curl_term =
                        cem::material::contract_vectors(curl[i][qp], curl[j][qp], inv_mu, dim) *
                        JxW[qp];
                    const double mass_eps = (omega * omega) *
                                            cem::material::contract_vectors(
                                                phi_v[i][qp], phi_v[j][qp], epsl_tensor, dim) *
                                            JxW[qp];
                    const double mass_sig = (omega)*cem::material::contract_vectors(
                                                phi_v[i][qp], phi_v[j][qp], sigma_tensor, dim) *
                                            JxW[qp];

                    const double real_part = curl_term - mass_eps;
                    const double imag_part = -mass_sig; // -i·ω·σ
                    ke(i, j) += Number(std::complex<double>(real_part, imag_part));
                }
                // 体 RHS 在此不构造
            }
        }

        // 边界项装配：遍历该单元的所有边界侧
        const BoundaryInfo & binfo = mesh.get_boundary_info();
        for (unsigned int s = 0; s < elem->n_sides(); ++s)
        {
            std::vector<boundary_id_type> side_ids;
            binfo.boundary_ids(elem, s, side_ids);
            if (side_ids.empty())
                continue; // 非边界侧

            // 初始化面 FE
            fe_face->reinit(elem, s);
            const auto & JxWf = fe_face->get_JxW();
            const auto & phif_v = fe_face->get_phi();
            const auto & normals = fe_face->get_normals();
            const auto & qpoints = fe_face->get_xyz();

            // 该侧可能关联多个边界 id：逐一处理
            for (const auto & bid : side_ids)
            {
                // 查找该边界 id 对应的规格，不存在则跳过
                auto itb = bnd_spec_map.find(static_cast<int>(bid));
                if (itb == bnd_spec_map.end())
                    continue;

                const cem::BoundarySpec & spec = itb->second;
                std::unique_ptr<BoundaryConditionObject> bc_obj =
                    make_boundary_condition_object(spec.type);
                if (!bc_obj)
                    continue;

                // 使用本地占位矩阵，避免空指针伪引用导致未定义行为
                DenseMatrix<Number> ke_dummy;
                bc_obj->assemble_face(spec,
                                      system,
                                      phif_v,
                                      JxWf,
                                      normals,
                                      qpoints,
                                      n,
                                      ke_dummy,
                                      fe_local);

                if (spec.type == BoundaryType::Dirichlet)
                {
                    std::unique_ptr<BoundaryConditionObject> dirichlet_obj =
                        make_boundary_condition_object(BoundaryType::Dirichlet);
                    // 使用本地空集合占位 rows，RHS-only 阶段不进行 zero_rows
                    std::unordered_set<libMesh::dof_id_type> rows_dummy;
                    dirichlet_obj->accumulate_dirichlet_face(spec,
                                                             system,
                                                             phif_v,
                                                             JxWf,
                                                             normals,
                                                             qpoints,
                                                             n,
                                                             dofs,
                                                             rows_dummy,
                                                             dirichlet_value_num,
                                                             dirichlet_value_den);
                }
            }
        }

        // 将局部矩阵/向量散装到全局
        dof_map.constrain_element_matrix_and_vector(ke, fe_local, dofs);
        sys.matrix->add_matrix(ke, dofs);
        sys.rhs->add_vector(fe_local, dofs);
    }

    // 在进行行清零之前，先装配完成并关闭矩阵（PETSc 的 MatZeroRows 需要已装配矩阵）
    sys.matrix->close();

    // 强制施加 Dirichlet：将对应行置零并在对角填 1，同时 RHS 置为投影值（复数）
    // 注意：zero_rows 是矩阵上的集体操作，必须在所有进程上调用；即使本地没有需要处理的行，
    // 也应以空列表参与调用以避免并行死锁。
    {
        // 计算各行目标值：value = numer / den（复数/实数）
        std::unordered_map<libMesh::dof_id_type, std::complex<double>> dirichlet_value;
        for (const auto & kv : dirichlet_value_den)
        {
            const auto row = kv.first;
            const double den = kv.second;
            std::complex<double> val(0.0, 0.0);
            auto itn = dirichlet_value_num.find(row);
            if (itn != dirichlet_value_num.end() && std::abs(den) > 0.0)
                val = itn->second / den;
            dirichlet_value[row] = val;
        }

        // 将需要保留的行填 1，其余设为 0（全员调用 zero_rows；允许空列表）
        std::vector<libMesh::numeric_index_type> dirichlet_rows_ordered;
        dirichlet_rows_ordered.reserve(dirichlet_rows_keep.size());
        for (const auto & rid : dirichlet_rows_keep)
            dirichlet_rows_ordered.push_back(static_cast<libMesh::numeric_index_type>(rid));
        std::sort(dirichlet_rows_ordered.begin(), dirichlet_rows_ordered.end());
        sys.matrix->zero_rows(dirichlet_rows_ordered, 1.0);

        // 设置 RHS：一次性写入（复数）
        for (const auto & kv : dirichlet_value)
            sys.rhs->set(kv.first, Number(kv.second));
    }

    sys.rhs->close();
}

// 新增方法: 仅根据外部给定的边界规格重装 RHS，不改动系统矩阵
// 适用场景: 双极化计算中复用矩阵，仅切换 RHS 进行两次回代
void
Assembler::assemble_rhs_with_specs(ElectromagneticsSystem & system,
                                   const std::vector<BoundarySpec> & specs)
{
    using namespace libMesh;

    auto & es = system.equation_systems();
    auto & sys = system.linear_system();

    const MeshBase & mesh = es.get_mesh();
    const unsigned int dim = mesh.mesh_dimension();

    if (mesh.n_active_elem() == 0 || dim == 0)
    {
        sys.rhs->zero();
        sys.rhs->close();
        return;
    }

    const DofMap & dof_map = sys.get_dof_map();
    const unsigned int var = system.var_index();
    const FEType fe_type = sys.variable_type(var);

    std::unique_ptr<FEVectorBase> fe_face = FEVectorBase::build(dim, fe_type);
    QGauss qface(dim > 0 ? dim - 1 : 0, SECOND);
    fe_face->attach_quadrature_rule(&qface);

    DenseVector<Number> fe_local;
    std::vector<dof_id_type> dofs;

    // 将边界规格整理为 id->spec 的快速查询表
    std::unordered_map<int, cem::BoundarySpec> bnd_spec_map;
    for (const auto & sp : specs)
        bnd_spec_map[sp.boundary_id] = sp;

    // 非齐次切向 Dirichlet 的 L2 投影累计（按“行 id”累计复数 numer 与公共分母）
    std::unordered_map<libMesh::dof_id_type, std::complex<double>> dirichlet_value_num;
    std::unordered_map<libMesh::dof_id_type, double> dirichlet_value_den;

    // 清零 RHS
    sys.rhs->zero();

    for (auto it = mesh.active_local_elements_begin(); it != mesh.active_local_elements_end(); ++it)
    {
        const Elem * elem = *it;

        dof_map.dof_indices(elem, dofs, var);
        const unsigned int n = dofs.size();

        fe_local.resize(n);
        fe_local.zero();

        const BoundaryInfo & binfo = mesh.get_boundary_info();
        for (unsigned int s = 0; s < elem->n_sides(); ++s)
        {
            std::vector<boundary_id_type> side_ids;
            binfo.boundary_ids(elem, s, side_ids);
            if (side_ids.empty())
                continue;

            fe_face->reinit(elem, s);
            const auto & JxWf = fe_face->get_JxW();
            const auto & phif_v = fe_face->get_phi();
            const auto & normals = fe_face->get_normals();
            const auto & qpoints = fe_face->get_xyz();

            for (const auto & bid : side_ids)
            {
                auto itb = bnd_spec_map.find(static_cast<int>(bid));
                if (itb == bnd_spec_map.end())
                    continue;

                const cem::BoundarySpec & spec = itb->second;
                std::unique_ptr<BoundaryConditionObject> bc_obj =
                    make_boundary_condition_object(spec.type);
                if (!bc_obj)
                    continue;

                // 使用本地占位矩阵，避免空指针伪引用导致未定义行为
                DenseMatrix<Number> ke_dummy;
                bc_obj->assemble_face(spec,
                                      system,
                                      phif_v,
                                      JxWf,
                                      normals,
                                      qpoints,
                                      n,
                                      ke_dummy,
                                      fe_local);

                if (spec.type == BoundaryType::Dirichlet)
                {
                    std::unique_ptr<BoundaryConditionObject> dirichlet_obj =
                        make_boundary_condition_object(BoundaryType::Dirichlet);
                    // 使用本地空集合占位 rows，RHS-only 阶段不进行 zero_rows
                    std::unordered_set<libMesh::dof_id_type> rows_dummy;
                    dirichlet_obj->accumulate_dirichlet_face(spec,
                                                             system,
                                                             phif_v,
                                                             JxWf,
                                                             normals,
                                                             qpoints,
                                                             n,
                                                             dofs,
                                                             rows_dummy,
                                                             dirichlet_value_num,
                                                             dirichlet_value_den);
                }
            }
        }

        dof_map.constrain_element_vector(fe_local, dofs);
        sys.rhs->add_vector(fe_local, dofs);
    }
    
    // 写入 Dirichlet 投影 RHS（不进行 zero_rows）
    for (const auto & kv : dirichlet_value_den)
    {
        const auto row = kv.first;
        const double den = kv.second;
        std::complex<double> val(0.0, 0.0);
        auto itn = dirichlet_value_num.find(row);
        if (itn != dirichlet_value_num.end() && std::abs(den) > 0.0)
            val = itn->second / den;
        sys.rhs->set(row, Number(val));
    }

    sys.rhs->close();
}

} // namespace cem